package com.starhub.integration.sdk.custom.core;

import lombok.Builder;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.RequestBody;
import okhttp3.Response;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.starhub.utils.Util;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 私有化重排序模型客户端
 * 支持基于xinfrence部署的重排序模型
 * @version: V1.0
 * @author: liuhf
 * @time 2025-08-08
 */
@Builder
public class CustomRerankerClient {
    private OkHttpClient okHttpClient;
    private String apiHost;
    private String modelName;
    private String apiPath;

    /**
     * 使用私有化重排序模型对文本列表进行评分
     * 
     * @param texts 需要评分的文本列表
     * @param query 查询文本
     * @return 评分列表
     * @throws IOException 当API调用失败或响应解析失败时抛出
     */
    public List<Double> getRerankerScores(List<String> texts, String query) throws IOException {
        if (texts == null || texts.isEmpty() || query == null || query.trim().isEmpty()) {
            throw new IOException("输入文本或查询不能为空");
        }

        String url = apiHost + apiPath;

        Map<String, Object> requestParams = new HashMap<>();
        requestParams.put("query", query);
        requestParams.put("documents", texts);
        if (modelName != null && !modelName.trim().isEmpty()) {
            requestParams.put("model", modelName);
        }

        RequestBody body = RequestBody.create(
                JSON.toJSONString(requestParams),
                MediaType.parse("application/json; charset=utf-8")
        );

        Request.Builder requestBuilder = new Request.Builder()
                .url(url)
                .addHeader("Content-Type", "application/json")
                .post(body);

        // 如果配置了API Key，添加到请求头
        /*  if (token != null && !token.trim().isEmpty()) {
            requestBuilder.addHeader("Authorization", "Bearer " + token);
        }*/

        Request request = requestBuilder.build();

        try (Response response = okHttpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException(String.format("API请求失败，状态码: %d", response.code()));
            }

            String responseBody = response.body().string();
            if (Util.isEmpty(responseBody)) {
                throw new IOException("API返回响应为空");
            }

            return parseRerankerResponse(responseBody);
        } catch (IOException e) {
            throw new IOException("调用私有化重排序API失败: " + e.getMessage(), e);
        }
    }

    /**
     * 解析私有化重排序模型的响应
     * 
     * @param responseBody API响应内容
     * @return 评分列表
     * @throws IOException 当响应解析失败时抛出
     */
    private List<Double> parseRerankerResponse(String responseBody) throws IOException {
        try {
            JSONObject jsonObject = JSON.parseObject(responseBody);

            // 检查错误信息
            String error = jsonObject.getString("error");
            if (Util.isNotEmpty(error)) {
                throw new IOException("API返回错误: " + error);
            }

            // 尝试多种响应格式
            List<Double> scores = new ArrayList<>();

            // 格式1: OpenAI兼容格式
            JSONArray dataArray = jsonObject.getJSONArray("data");
            if (dataArray != null && !dataArray.isEmpty()) {
                for (int i = 0; i < dataArray.size(); i++) {
                    JSONObject data = dataArray.getJSONObject(i);
                    Double score = data.getDouble("score");
                    if (score != null) {
                        scores.add(score);
                    }
                }
                if (!scores.isEmpty()) {
                    return scores;
                }
            }

            // 格式2: 自定义格式 - results数组
            JSONArray resultsArray = jsonObject.getJSONArray("results");
            if (resultsArray != null && !resultsArray.isEmpty()) {
                for (int i = 0; i < resultsArray.size(); i++) {
                    JSONObject result = resultsArray.getJSONObject(i);
                    Double score = result.getDouble("score");
                    if (score != null) {
                        scores.add(score);
                    }
                }
                if (!scores.isEmpty()) {
                    return scores;
                }
            }

            // 格式3: 直接返回scores数组
            JSONArray scoresArray = jsonObject.getJSONArray("scores");
            if (scoresArray != null && !scoresArray.isEmpty()) {
                for (int i = 0; i < scoresArray.size(); i++) {
                    Object scoreObj = scoresArray.get(i);
                    if (scoreObj instanceof Number) {
                        scores.add(((Number) scoreObj).doubleValue());
                    }
                }
                if (!scores.isEmpty()) {
                    return scores;
                }
            }

            // 格式4: 单个score字段
            Double singleScore = jsonObject.getDouble("score");
            if (singleScore != null) {
                scores.add(singleScore);
                return scores;
            }

            throw new IOException("无法解析API响应格式，响应内容: " + responseBody);

        } catch (Exception e) {
            throw new IOException("解析API响应失败: " + e.getMessage(), e);
        }
    }
} 